// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/aggregates/approx_most_frequent_aggregate.h>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/simple_aggregate_adapter.h>
#include <pollux/exec/strings.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/lib/approx_most_frequent_stream_summary.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        template<typename T>
        struct Accumulator {
            functions::ApproxMostFrequentStreamSummary<T, AlignedStlAllocator<T, 16> >
            summary;

            explicit Accumulator(HashStringAllocator *allocator)
                : summary(AlignedStlAllocator<T, 16>(allocator)) {
            }

            void insert(T value, int64_t count = 1) {
                summary.insert(value, count);
            }
        };

        template<>
        struct Accumulator<StringView> {
            functions::ApproxMostFrequentStreamSummary<
                StringView,
                AlignedStlAllocator<StringView, 16> >
            summary;
            HashStringAllocator *allocator;
            Strings strings;

            explicit Accumulator(HashStringAllocator *allocator)
                : summary(AlignedStlAllocator<StringView, 16>(allocator)),
                  allocator(allocator) {
            }

            ~Accumulator() {
                strings.free(*allocator);
            }

            void insert(StringView value, int64_t count = 1) {
                if (!value.isInline() && !summary.contains(value)) {
                    value = strings.append(value, *allocator);
                }
                summary.insert(value, count);
            }
        };

        template<typename T>
        struct ApproxMostFrequentAggregate : exec::Aggregate {
            explicit ApproxMostFrequentAggregate(const TypePtr &resultType)
                : Aggregate(resultType) {
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(Accumulator<T>);
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool) override {
                decodeArguments(rows, args);
                rows.applyToSelected([&](auto row) {
                    if (!decodedValues_.is_null_at(row)) {
                        auto *accumulator = initAccumulator(groups[row]);
                        accumulator->insert(decodedValues_.value_at<T>(row));
                    }
                });
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool) override {
                addIntermediate<false>(groups, rows, args);
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool) override {
                decodeArguments(rows, args);
                auto *accumulator = initAccumulator(group);
                rows.applyToSelected([&](auto row) {
                    if (!decodedValues_.is_null_at(row)) {
                        accumulator->insert(decodedValues_.value_at<T>(row));
                    }
                });
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool) override {
                addIntermediate<true>(group, rows, args);
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                (*result)->resize(numGroups);
                if (buckets_ == kMissingArgument) {
                    // No data has been added.
                    for (int i = 0; i < numGroups; ++i) {
                        POLLUX_DCHECK_EQ(value<Accumulator<T>>(groups[i])->summary.size(), 0);
                        (*result)->set_null(i, true);
                    }
                    return;
                }
                POLLUX_USER_CHECK_LE(buckets_, std::numeric_limits<int>::max());
                auto map_vector = (*result)->as<MapVector>();
                auto [keys, values] = prepareFinalResult(groups, numGroups, map_vector);
                vector_size_t entryCount = 0;
                for (int i = 0; i < numGroups; ++i) {
                    auto *summary = &value<Accumulator<T> >(groups[i])->summary;
                    int size = std::min<int>(buckets_, summary->size());
                    if (size == 0) {
                        map_vector->set_null(i, true);
                    } else {
                        summary->topK(
                            buckets_,
                            keys->mutableRawValues() + entryCount,
                            values->mutableRawValues() + entryCount);
                        if constexpr (std::is_same_v<T, StringView>) {
                            // Populate the string buffers.
                            for (int j = 0; j < size; ++j) {
                                keys->set(entryCount + j, keys->valueAtFast(entryCount + j));
                            }
                        }
                        entryCount += size;
                    }
                    map_vector->setOffsetAndSize(i, entryCount - size, size);
                }
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto rowVec = (*result)->as<RowVector>();
                POLLUX_CHECK(rowVec);
                rowVec->childAt(0) = std::make_shared<ConstantVector<int64_t> >(
                    rowVec->pool(),
                    numGroups,
                    false,
                    BIGINT(),
                    static_cast<int64_t &&>(buckets_));
                rowVec->childAt(1) = std::make_shared<ConstantVector<int64_t> >(
                    rowVec->pool(),
                    numGroups,
                    false,
                    BIGINT(),
                    static_cast<int64_t &&>(capacity_));
                auto values = rowVec->childAt(2)->as<ArrayVector>();
                auto counts = rowVec->childAt(3)->as<ArrayVector>();
                rowVec->resize(numGroups);
                values->resize(numGroups);
                counts->resize(numGroups);

                auto v = values->elements()->template as_flat_vector<T>();
                auto c = counts->elements()->template as_flat_vector<int64_t>();
                vector_size_t entryCount = 0;
                for (int i = 0; i < numGroups; ++i) {
                    auto *accumulator = value<const Accumulator<T>>(groups[i]);
                    entryCount += accumulator->summary.size();
                }
                v->resize(entryCount);
                c->resize(entryCount);
                v->reset_nulls();
                c->reset_nulls();

                entryCount = 0;
                for (int i = 0; i < numGroups; ++i) {
                    auto *summary = &value<const Accumulator<T>>(groups[i])->summary;
                    if (summary->size() == 0) {
                        rowVec->set_null(i, true);
                    } else {
                        if constexpr (std::is_same_v<T, StringView>) {
                            for (int j = 0; j < summary->size(); ++j) {
                                v->set(entryCount + j, summary->values()[j]);
                            }
                        } else {
                            memcpy(
                                v->mutableRawValues() + entryCount,
                                summary->values(),
                                sizeof(T) * summary->size());
                        }
                        memcpy(
                            c->mutableRawValues() + entryCount,
                            summary->counts(),
                            sizeof(int64_t) * summary->size());
                        values->setOffsetAndSize(i, entryCount, summary->size());
                        counts->setOffsetAndSize(i, entryCount, summary->size());
                        entryCount += summary->size();
                    }
                }
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                for (auto index: indices) {
                    new(groups[index] + offset_) Accumulator<T>(allocator_);
                }
            }

            void destroyInternal(melon::Range<char **> groups) override {
                destroyAccumulators<Accumulator<T> >(groups);
            }

        private:
            void decodeArguments(
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args) {
                POLLUX_CHECK_EQ(args.size(), 3);
                DecodedVector decodedBuckets(*args[0], rows);
                decodedValues_.decode(*args[1], rows);
                DecodedVector decodedCapacity(*args[2], rows);
                setConstantArgument("Buckets", buckets_, decodedBuckets);
                setConstantArgument("Capacity", capacity_, decodedCapacity);
            }

            static void
            setConstantArgument(const char *name, int64_t &val, int64_t newVal) {
                POLLUX_USER_CHECK_GT(newVal, 0, "{} must be positive", name);
                if (val == kMissingArgument) {
                    val = newVal;
                } else {
                    POLLUX_USER_CHECK_EQ(
                        newVal, val, "{} argument must be constant for all input rows", name);
                }
            }

            static void setConstantArgument(
                const char *name,
                int64_t &val,
                const DecodedVector &vec) {
                POLLUX_USER_CHECK(
                    vec.isConstantMapping(),
                    "{} argument must be constant for all input rows",
                    name);
                setConstantArgument(name, val, vec.value_at<int64_t>(0));
            }

            Accumulator<T> *initAccumulator(char *group) {
                auto accumulator = value<Accumulator<T> >(group);
                POLLUX_USER_CHECK_LE(capacity_, std::numeric_limits<int>::max());
                accumulator->summary.setCapacity(capacity_);
                return accumulator;
            }

            template<bool kSingleGroup>
            void addIntermediate(
                std::conditional_t<kSingleGroup, char *, char **> group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args) {
                POLLUX_CHECK_EQ(args.size(), 1);
                DecodedVector decoded(*args[0], rows);
                auto rowVec = static_cast<const RowVector *>(decoded.base());
                auto buckets = rowVec->childAt(0)->as<SimpleVector<int64_t> >();
                auto capacity = rowVec->childAt(1)->as<SimpleVector<int64_t> >();
                auto values = rowVec->childAt(2)->as<ArrayVector>();
                auto counts = rowVec->childAt(3)->as<ArrayVector>();
                POLLUX_CHECK(buckets);
                POLLUX_CHECK(capacity);
                POLLUX_CHECK(values);
                POLLUX_CHECK(counts);

                auto v = values->elements()->template as_flat_vector<T>();
                auto c = counts->elements()->template as_flat_vector<int64_t>();
                POLLUX_CHECK(v);
                POLLUX_CHECK(c);

                Accumulator<T> *accumulator = nullptr;
                rows.applyToSelected([&](auto row) {
                    if (decoded.is_null_at(row)) {
                        return;
                    }
                    int i = decoded.index(row);
                    setConstantArgument("Buckets", buckets_, buckets->value_at(i));
                    setConstantArgument("Capacity", capacity_, capacity->value_at(i));
                    if constexpr (kSingleGroup) {
                        if (!accumulator) {
                            accumulator = initAccumulator(group);
                        }
                    } else {
                        accumulator = initAccumulator(group[row]);
                    }
                    auto size = values->sizeAt(i);
                    POLLUX_DCHECK_EQ(counts->sizeAt(i), size);
                    auto vo = values->offsetAt(i);
                    auto co = counts->offsetAt(i);
                    for (int j = 0; j < size; ++j) {
                        accumulator->insert(v->value_at(vo + j), c->value_at(co + j));
                    }
                });
            }

            std::pair<FlatVector<T> *, FlatVector<int64_t> *>
            prepareFinalResult(char **groups, int32_t numGroups, MapVector *result) {
                POLLUX_CHECK(result);
                auto keys = result->mapKeys()->as_unchecked<FlatVector<T> >();
                auto values = result->mapValues()->as_unchecked<FlatVector<int64_t> >();
                POLLUX_CHECK(keys);
                POLLUX_CHECK(values);
                vector_size_t entryCount = 0;
                for (int i = 0; i < numGroups; ++i) {
                    auto *summary = &value<const Accumulator<T>>(groups[i])->summary;
                    entryCount += std::min<int>(buckets_, summary->size());
                }
                keys->resize(entryCount);
                values->resize(entryCount);
                return std::make_pair(keys, values);
            }

            static constexpr int64_t kMissingArgument = -1;
            DecodedVector decodedValues_;
            int64_t buckets_ = kMissingArgument;
            int64_t capacity_ = kMissingArgument;
        };

        class ApproxMostFrequentBooleanAggregate {
        public:
            using InputType =
            Row</*buckets*/ int64_t, /*value*/ bool, /*capacity*/ int64_t>;

            using IntermediateType =
            Row</*buckets*/ int64_t,
                /*capacity*/ int64_t,
                /*values*/ Array<bool>,
                /*counts*/ Array<int64_t> >;

            using OutputType = Map<bool, int64_t>;

            static bool toIntermediate(
                exec::out_type<IntermediateType> &out,
                int64_t buckets,
                bool value,
                int64_t capacity) {
                out.get_writer_at<0>() = buckets;
                out.get_writer_at<1>() = capacity;

                auto &valuesWriter = out.get_writer_at<2>();
                valuesWriter.add_item() = true;
                valuesWriter.add_item() = false;

                auto &countsWriter = out.get_writer_at<3>();
                countsWriter.add_item() = value ? 1 : 0;
                countsWriter.add_item() = value ? 0 : 1;

                return true;
            }

            struct AccumulatorType {
                int64_t numTrue{0};
                int64_t numFalse{0};

                explicit AccumulatorType(
                    HashStringAllocator * /*allocator*/,
                    ApproxMostFrequentBooleanAggregate * /*fn*/) {
                }

                void addInput(
                    HashStringAllocator * /*allocator*/,
                    int64_t /*buckets*/,
                    bool value,
                    int64_t /*capacity*/) {
                    if (value) {
                        ++numTrue;
                    } else {
                        ++numFalse;
                    }
                }

                void combine(
                    HashStringAllocator * /*allocator*/,
                    exec::arg_type<IntermediateType> other) {
                    POLLUX_CHECK(other.at<2>().has_value());
                    POLLUX_CHECK(other.at<3>().has_value());

                    const auto &values = *other.at<2>();
                    POLLUX_CHECK_EQ(2, values.size());

                    POLLUX_CHECK_EQ(values[0].value(), true);
                    POLLUX_CHECK_EQ(values[1].value(), false);

                    const auto &counts = *other.at<3>();
                    POLLUX_CHECK_EQ(2, counts.size());

                    numTrue += counts[0].value();
                    numFalse += counts[1].value();
                }

                bool writeFinalResult(exec::out_type<OutputType> &out) {
                    if (numTrue > 0) {
                        auto [keyWriter, valueWriter] = out.add_item();
                        keyWriter = true;
                        valueWriter = numTrue;
                    }

                    if (numFalse > 0) {
                        auto [keyWriter, valueWriter] = out.add_item();
                        keyWriter = false;
                        valueWriter = numFalse;
                    }

                    return true;
                }

                bool writeIntermediateResult(exec::out_type<IntermediateType> &out) {
                    // Write some hard-coded values for 'buckets' and 'capacity'. These are
                    // not used.
                    out.get_writer_at<0>() = 2;
                    out.get_writer_at<1>() = 2;

                    auto &valuesWriter = out.get_writer_at<2>();
                    valuesWriter.add_item() = true;
                    valuesWriter.add_item() = false;

                    auto &countsWriter = out.get_writer_at<3>();
                    countsWriter.add_item() = numTrue;
                    countsWriter.add_item() = numFalse;

                    return true;
                }
            };
        };

        template<TypeKind kKind>
        std::unique_ptr<exec::Aggregate> makeApproxMostFrequentAggregate(
            const std::string &name,
            core::AggregationNode::Step step,
            const std::vector<TypePtr> &argTypes,
            const TypePtr &resultType,
            const TypePtr &valueType) {
            if constexpr (
                kKind == TypeKind::TINYINT || kKind == TypeKind::SMALLINT ||
                kKind == TypeKind::INTEGER || kKind == TypeKind::BIGINT ||
                kKind == TypeKind::VARCHAR) {
                return std::make_unique<
                    ApproxMostFrequentAggregate<typename TypeTraits<kKind>::NativeType> >(
                    resultType);
            }

            if (kKind == TypeKind::BOOLEAN) {
                return std::make_unique<
                    exec::SimpleAggregateAdapter<ApproxMostFrequentBooleanAggregate> >(
                    step, argTypes, resultType);
            }

            POLLUX_USER_FAIL(
                "Unsupported value type for {} aggregation {}",
                name,
                valueType->toString());
        }
    } // namespace

    void registerApproxMostFrequentAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;
        for (const auto &valueType:
             {
                 "boolean",
                 "tinyint",
                 "smallint",
                 "integer",
                 "bigint",
                 "varchar",
                 "json"
             }) {
            signatures.push_back(
                exec::AggregateFunctionSignatureBuilder()
                .returnType(fmt::format("map({},bigint)", valueType))
                .intermediateType(fmt::format(
                    "row(bigint, bigint, array({}), array(bigint))", valueType))
                .argumentType("bigint")
                .argumentType(valueType)
                .argumentType("bigint")
                .build());
        }
        auto name = prefix + kApproxMostFrequent;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                auto &valueType = exec::isPartialOutput(step)
                                      ? resultType->childAt(2)->childAt(0)
                                      : resultType->childAt(0);
                return POLLUX_DYNAMIC_TYPE_DISPATCH(
                            makeApproxMostFrequentAggregate,
                            valueType->kind(),
                            name,
                            step,
                            argTypes,
                            resultType,
                            valueType);
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
